# FIXME -- I realized that there are six possibilities to compute
# shapes: You choose spherical or ellipsoidal window, then choose
# among no weighting, spherical weighting, and ellipsoidal weighting.

# FIXME -- Sauron plots use old Reffs, regenerate them at some point.

import os, types

from numpy import argsort, array

import sauron, plots, util, local, snapshot, properties

def hasPlot(fn):    
    return os.path.isfile(fn) or os.path.isfile(fn + '.gz')

def allPlots(qty, verbose=True, dataPrefix=None, plotPrefix=None, dryRun=False):
    if plotPrefix is None: plotPrefix = 'Plots'
    if dataPrefix is None: dataPrefix = 'data'
    if verbose: print '*** Plots ***'
    summaryPlots(qty, plotPrefix=plotPrefix, verbose=verbose, dryRun=dryRun)

    for sim in qty.keys():
        simDataPrefix = dataPrefix + os.sep + sim
        simPlotPrefix = plotPrefix + os.sep + sim
        simulationPlots(simDataPrefix, qty[sim],
                        plotPrefix=simPlotPrefix, verbose=verbose, dryRun=dryRun)

        for snap in util.dictKeys(qty[sim]):
            snapPlotPrefix = plotPrefix + os.sep + sim + os.sep + snap
            dataFn = dataPrefix + os.sep + sim + os.sep + snap
            snapshotPlots(dataFn, qty[sim][snap],
                          plotPrefix=snapPlotPrefix, verbose=verbose, dryRun=dryRun)

def summaryPlots(qty, plotPrefix=None, verbose=True, dryRun=False):
    if plotPrefix is None: plotPrefix = '.'
    outputPrefix = ''

    for f in summaryPlotList:
        plotFn = plotPrefix + os.sep + f.func_name + '.eps'
        if verbose: print outputPrefix + f.func_name + '...', 
        if hasPlot(plotFn):
            if verbose: print 'up to date'
        elif dryRun:
            if verbose: print 'would plot'            
        else: 
            if verbose: print 'plotting'
            f(qty, filename=plotFn)

def simulationPlots(dataPrefix, sim, plotPrefix=None, verbose=True, dryRun=False):    
    if plotPrefix is None: plotPrefix = '.'
    outputPrefix = ''
    if verbose: print outputPrefix + sim['simulation'] + '...',

    plotFns = [plotPrefix + os.sep + f.func_name + '.eps'
               for f in simulationPlotList]

    snaps = properties.timeOrderedSnapshots(sim)
    dataFns = [dataPrefix + os.sep + snap for snap in snaps]
    filesNotFound = [fn for fn in dataFns if not os.path.isfile(fn)]

    if filesNotFound: 
        # TODO -- not necessarily a fatal error, can just ignore the
        # files not found
        print outputPrefix + "Warning, file not found:", filesNotFound
        return
    
    if util.every([ hasPlot(fn) for fn in plotFns]):
        if verbose: print "up to date"
        return
    
    if verbose: print "reading"
    dataDict = {}
    # FIXME -- don't actually load data for all snapshots
    #         for snap in snaps:
    #             dataFn = dataPrefix + os.sep + snap
    #             dataDict[snap] = snapshot.GadgetSnapshot(dataFn)
    #             dataDict[snap].toBodyCoordinates(props)
    for f, plotFn in zip(simulationPlotList, plotFns):
        if verbose: print outputPrefix + '  ' + f.__name__ + '...',
        if hasPlot(plotFn):
            if verbose: print 'up to date'
        elif dryRun:
            if verbose: print 'would plot'
        else: 
            if verbose: print 'plotting'
            os.system("touch " + plotFn)
            try: f(dataDict, sim, filename = plotFn)
            except (KeyboardInterrupt, SystemExit):
                os.remove(plotFn)
            except Exception:
                print "FAILED to plot", plotFn
                os.remove(plotFn)
                if debug: raise
            

        
def snapshotPlots(dataFn, props, plotPrefix=None, verbose=True,
                  debug=False, dryRun=False):
    if plotPrefix is None: plotPrefix = '.'
    outputPrefix = '  '
    if verbose: print outputPrefix + props['snapshot'] + '...',
    if not os.path.isfile(dataFn):
        print outputPrefix + "Warning, file not found:", dataFn
        return

    plotFns = [ plotPrefix + os.sep + f.func_name + '.eps'
                for f in snapshotPlotList]    
    # Check if any of the plots need updated    
    if util.every([ hasPlot(fn) for fn in plotFns]):
        if verbose: print 'up to date'
        return 

    if verbose: print 'reading'
    data = snapshot.GadgetSnapshot(dataFn)
    data.toBodyCoordinates(props)
    
    for f, plotFn in zip(snapshotPlotList, plotFns):
        if verbose: print outputPrefix + '  ' + f.func_name, 
        # Lock the plot
        if hasPlot(plotFn):
            if verbose: print 'up to date'
        elif dryRun:
            if verbose: print 'would plot'
        else:
            if verbose: print 'plotting'
            os.system("touch " + plotFn)
            try: f(data, props, filename=plotFn)
            except (KeyboardInterrupt, SystemExit):
                os.remove(plotFn)
            except Exception:
                print "FAILED to plot", plotFn
                os.remove(plotFn)
                if debug: raise
                
            
##################################################
## Rickety code
##################################################

def AAA_plot_pages(qty,filename=None):
    # Don't do anything if filename not specified
    if not filename: return
    
    # filename serves as a prefix - strip .eps if necessary
    filename = filename.replace('.eps','')
    
    ### snapshot selectors
    def last_snap(lst):
        if lst: return [lst[-1]]
        return []
    def all_snaps(lst): return lst
    def same_time_snaps(lst):
        if len(lst) == 1: return lst
        else: return ['snapshot_048']

    classes = [ # ('suffix', snapshot_selector, [list, of, simnames])
        ('', last_snap, util.sorted(qty.keys())),
        ('_all', all_snaps, util.sorted(qty.keys())),
        ('_mass', last_snap,
         ['G0G0-u1', 'G0G0a-u1','G0G0r-u1b','G1G1a-u1','G1G1r-u1', 'G2r',
          'G3', 'G3b', 'G3gf1G3gf1b-u1', 'G3gf2G3gf2b-u1', 'G3il-u1', 'G3r']),
        ('_feedback', last_snap,
         ['Sbc201a-u50', 'Sbc201a-u51', 'Sbc201a-u53', 'Sbc201a-u52',
          'Sbc201a-u4', 'Sbc201a-u43']),
        ('_minor',last_snap,
         ['G1G0-u3', 'G1G0r-u3', 'G2G0', 'G2G0r', 'G2G1', 'G2G1r',
          'G3G0e-u3', 'G3G1', 'G3G1-u3', 'G3G1r', 'G3G2',
          'G3G2-u3-rerun', 'G3G2a-u3', 'G3G2b-u3', 'G3G2c-u3',
          'G3G2d-u3', 'G3G2e-u3', 'G3G2f-u3', 'G3G2g-u3', 'G3G2h-u3',
          'G3G2r']),
        ('_orbit',last_snap,         
         ['Sbc201a-u4', 'Sbc202-u4', 'Sbc203-u4', 'Sbc204-u4',
          'Sbc205-u4', 'Sbc206-u4', 'Sbc207-u4', 'Sbc208-u4',
          'Sbc209-u4', 'Sbc210-u4', 'Sbc211-u4', 'Sbc212-u4',
          'Sbc213-u4', 'Sbc214-u4', 'Sbc215-u4', 'Sbc216-u4',
          'Sbc217-u4', 'Sbc218-u4', 'Sbc219-u4']),
        ('_resolution',same_time_snaps,         
         ['Sbc201a-u4', 'Sbc201a10x-n4a']),
        ]


    def write_html_file(filename, html_body):
        """Body is list of strings, one for each line"""
        
        html_header = """<!DOCTYPE HTML PUBLIC "-//IETF//DTD HTML//EN">
        <html>
        <head><title>Catalog</title></head>
        <body>
        """ 
        html_footer="</body></html>"
        f = open(filename,'w')
        f.writelines([html_header] + html_body + [html_footer])
        f.close()

    def snapshot_html_line(sim,snap,func):
        """all args are strings"""
        template = '<a href=../%s/%s><img border=0 src="../%s/%s/%s.png"></a>\n'
        return template % (sim,snap,sim,snap,func)


    for clas in classes:
        suffix, select, sims = clas
        # Simulation plots
        for func in simulationPlotList:

            body = ['<a href=../%s><img border=0 src="../%s/%s.png"></a>\n' %
                    (sim, sim ,func.func_name)
                    for sim in sims]                            
            write_html_file(filename+'/'+func.func_name+suffix+'.html', body)


        # Snapshot plots
        for func in snapshotPlotList:
            # Get the latest snapshot for each sim.        
            body = [ [ snapshot_html_line(sim,
                                          snap,
                                          func.func_name)
                       for snap in select(properties.timeOrderedSnapshots(qty[sim]))]
                     for sim in sims
                     if sim in qty]
            
            write_html_file(filename+'/'+func.func_name+suffix+'.html', util.flatten(body))        

def makeDirectoryTree(qty, prefix=None):
    if prefix is None: prefix = 'Plots'

    def makeDirectoryMaybe(path):
        if not os.path.isdir(path):
            print "Making ", path
            os.makedirs(path)
                    
    makeDirectoryMaybe(os.path.join(prefix, 'AAA_plot_pages'))
    for sim in util.dictKeys(qty):        
        for snap in util.dictKeys(qty[sim]):
            makeDirectoryMaybe(os.path.join(prefix, sim, snap))
        

def updateCatalog(qty=None, verbose=True, filename=None, dataPrefix=None, plotPrefix=None,
                  dryRun=False):
    if filename is None: filename = 'mergers.dat'
    if qty is None: qty = util.uncan(filename)
    
    makeDirectoryTree(qty, prefix=plotPrefix)
    allPlots(qty, verbose=verbose, dataPrefix=dataPrefix, plotPrefix=plotPrefix, dryRun=dryRun)
    if not dryRun:
        publishCatalog(verbose=verbose, prefix=plotPrefix)

def testUpdateCatalog(qty=None, dryRun=False):
    updateCatalog(qty=qty, verbose=True, filename='test.dat',
                  dataPrefix='test-data', plotPrefix='TestPlots', dryRun=dryRun)

def publishCatalog(verbose=True, prefix=None):
    if prefix is None: prefix = 'Plots'
    # Copy HTML files
    os.system(r"find ./"+prefix+"/ -type d -name snap\* -exec cp html/snapshot.html {}/index.html \;")
    os.system(r"find ./"+prefix+"/ -type d -not -name snap\* -exec cp html/simulation.html {}/AAA_index.html \;")
    os.system(r"find ./"+prefix+"/ -type d -exec cp html/png-makefile {}/Makefile \;")

    # futz a bit with the top level.
    os.system("cp html/summary.html ./"+prefix+"/AAA_index.html")
    
    # Convert image files
    os.system("find ./"+prefix+"/ -type d -exec make -s -C {} \;")

def updatePublicCatalog():
    cmd = 'rsync --archive --exclude=\*.gz Plots/* publichost:public_html/Catalog'
    os.system(cmd)
 
summaryPlotList = [AAA_plot_pages,
                   plots.kinemetryPaVelOff, 
                   plots.kinemetryPaVelMaj, 
                   plots.kinemetryPaVelMin, 
                   plots.kinemetryPaSigOff, 
                   plots.kinemetryPaSigMaj, 
                   plots.kinemetryPaSigMin, 
                   plots.kinemetryEpsVelOffPlot, 
                   plots.kinemetryEpsVelMinPlot, 
                   plots.kinemetryEpsVelMajPlot, 
                   plots.kinemetryEpsSigOffPlot, 
                   plots.kinemetryEpsSigMinPlot, 
                   plots.kinemetryEpsSigMajPlot, 
                   plots.kinemetryVOverSigmaOffPlot,
                   plots.kinemetryVOverSigmaMinPlot,
                   plots.kinemetryVOverSigmaMajPlot,
                   plots.kinemetryK5OK1Off,
                   plots.kinemetryK5OK1Maj,
                   plots.kinemetryK5OK1Min,
                   plots.kinemetryK4OK1Off,
                   plots.kinemetryK4OK1Maj,
                   plots.kinemetryK4OK1Min,
                   plots.sersicRmaxMin,
                   plots.sersicRminMin,
                   plots.sersicRmaxMaj, 
                   plots.sersicRminMaj,
                   ]
# FIXME
summaryPlotList = []    
#                      summary_table_plot,
#                      star_figure_rotation_times_plot,
#                      dark_figure_rotation_times_plot,
#                      dm_triaxiality_plot,
#                      stars_triaxiality_plot,
#                      v_over_sigma_histogram_plot,
#                      disk_fraction_histogram_plot,
#                      v_over_sigma_star_histogram_plot,
#                      v_over_sigma_vs_ellipticity_plot,
#                      v_over_sigma_star_vs_ellipticity_plot,
# #                     compare_spherical_unscaled_shapes,
# #                     compare_spherical_scaled_shapes,
# #                     compare_ellipsoid_scaled_shapes,
#                      galaxy_halo_angle_plot,
#                      dm_inner_vs_outer_shape_plot,
#                      star_inner_vs_outer_shape_plot,
#                      star_twist_summary_plot,
#                      dm_twist_summary_plot,
#                      star_triaxiality_parameter_plot,
#                      dark_triaxiality_parameter_plot,
                     
simulationPlotList = [
#     twist_vs_time_plot
    ]

snapshotPlotList = [plots.sphericalAnisotropyProfilePlot,
                    plots.cylindricalAnisotropyProfilePlot,
                    plots.sphericalSigmaProfilePlot,
                    plots.cylindricalSigmaProfilePlot,
                    plots.fastSauronMajPlot,
                    plots.fastSauronMinPlot,
                    plots.fastSauronOffPlot,
                    #plots.sauronMajPlot,
                    #plots.sauronMinPlot,
                    #plots.sauronOffPlot,
                    ]

#     sauron_min_plot,
#                       sauron_med_plot,
#                       sauron_maj_plot,
#                       sauron_min_med_plot,
#                       sauron_min_maj_plot,
#                       sauron_med_maj_plot,
#                       sauron_off_plot,
#                       pn_min_plot,
#                       pn_med_plot,
#                       pn_maj_plot,
#                       pn_min_med_plot,
#                       pn_min_maj_plot,
#                       pn_med_maj_plot,
#                       pn_off_plot,
#                       pyx_h4_vs_r_plot,
#                       major_v_over_sigma_plot,
#                       major_v_and_sigma_plot,
#                       minor_v_over_sigma_plot,
#                       minor_v_and_sigma_plot,
#                       face_on_surface_brightness_plot,
#                       edge_on_surface_brightness_plot,
#                       axis_with_radius_plot,
#                       axis_ratio_with_radius_plot,
#                       dm_axis_with_radius_plot,
#                       dm_axis_ratio_with_radius_plot,
#  #                     stars_axis_ratio_methods_plot,
#  #                     dm_axis_ratio_methods_plot,
#                       star_twist_vs_radius_plot,
#                       dm_twist_vs_radius_plot,
#                       basic_data_table_plot,
#                       density_plot,
#                       face_on_metallicity_image_plot,
#                       edge_on_metallicity_image_plot,
#                       edge_on_age_image_plot,
#                       face_on_age_image_plot]
